import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import welch

# --- 图表全局设置 ---
# 设置中文字体，以确保图表中的中文能正常显示。
# 'SimSun' 是宋体, 'Times New Roman' 是新罗马字体。
plt.rcParams['font.sans-serif'] = ['SimSun'] # 中文使用宋体
plt.rcParams['font.serif'] = ['Times New Roman'] # 英文和数字使用新罗马
plt.rcParams['font.family'] = ['serif', 'sans-serif'] # 优先使用衬线字体（新罗马），然后是无衬线字体（宋体）
# 解决在使用中文字体后，坐标轴负号'-'显示为方块的问题。
plt.rcParams['axes.unicode_minus'] = False

def find_displacement_data_path(root_path, region_processed_folder, flow_rate):
    """
    在指定的区域文件夹内，根据流速构建并查找 displacement_data.csv 文件的完整路径。

    参数:
        root_path (str): 数据存储的根目录。
        region_processed_folder (str): 对应测点区域的处理后文件夹全名 (例如, 'A BigF Jingu_processed')。
        flow_rate (float): 需要分析的流速值 (单位: m/s)。

    返回:
        str: displacement_data.csv 文件的完整路径，如果未找到则返回 None。
    """
    # 将浮点数格式的流速转换为特定格式的字符串，例如 3.0 -> "3.0"。
    flow_rate_str = f"{flow_rate:.1f}"
    # 根据您脚本中的命名规则，构建包含流速信息的文件夹名称。
    if flow_rate == 5.0 or flow_rate == 4.1: # 针对特殊流速命名进行适配
        flow_folder_name = f"1000f{flow_rate_str}ms"
    else:
        # 假设其他流速文件夹名称遵循您的示例，例如0.6 -> 1000f0.6ms
        flow_folder_name = f"1000f{flow_rate_str}ms"

    # 拼接成流速文件夹的完整路径。
    flow_folder_path = os.path.join(root_path, region_processed_folder, flow_folder_name)

    # 检查该路径是否存在且为文件夹。
    if os.path.isdir(flow_folder_path):
        # 获取下一级子文件夹，通常是 'C001H001S0001' 这样的名称。
        subfolders = [d for d in os.listdir(flow_folder_path) if os.path.isdir(os.path.join(flow_folder_path, d))]
        if subfolders:
            # 假设数据总是存在于找到的第一个子文件夹中。
            data_folder_path = os.path.join(flow_folder_path, subfolders[0])
            # 构建最终的 displacement_data.csv 文件路径。
            displacement_file_path = os.path.join(data_folder_path, 'displacement_data.csv')
            if os.path.exists(displacement_file_path):
                # 如果文件存在，返回其路径。
                return displacement_file_path
            else:
                # 如果文件不存在，打印提示信息。
                print(f"文件未找到: {displacement_file_path}")
                return None

    # 如果流速文件夹不存在，打印提示信息。
    print(f"在 {region_processed_folder} 中未找到流速为 {flow_rate} m/s 的文件夹 (预期路径: {flow_folder_path})")
    return None

def plot_psd_analysis(root_path, regions, flow_rates, output_folder):
    """
    对位移数据执行功率谱密度(PSD)分析，并绘制出紧凑美观的结果图表。

    参数:
        root_path (str): 数据存储的根目录。
        regions (list): 需要绘制的区域标识符列表 (例如, ['A', 'B', 'C'])。
        flow_rates (list): 需要分析的流速列表 (例如, [0.5, 1.2, 2.0, 3.0, 3.5])。
        output_folder (str): 保存输出图表的文件夹路径。
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # --- 修改点 1: 调整画布大小使其更高更窄，并取消X轴共享 ---
    # sharex=False 确保每个子图都有自己的X轴
    fig, axes = plt.subplots(len(regions), 1, figsize=(8, 10), sharex=False)
    # 如果只有一个子图, axes不是列表, 将其转换为列表以便统一处理
    if len(regions) == 1:
        axes = [axes]

    # 定义颜色和标记点样式，以模仿您提供的示例图风格。
    # 调整颜色以更好地匹配示例图
    colors = ['black', 'magenta', 'darkviolet', 'olive', 'steelblue', 'darkcyan']
    markers = ['s', 'v', 'o', 'p', 'D', '^']

    all_folders_in_root = os.listdir(root_path)

    for i, region_prefix in enumerate(regions):
        ax = axes[i] # 获取当前区域对应的子图对象。

        region_folder_name = None
        for folder in all_folders_in_root:
            if folder.startswith(region_prefix + ' ') and folder.endswith('_processed'):
                region_folder_name = folder
                break 

        if not region_folder_name:
            print(f"警告: 在 {root_path} 中未找到以前缀 '{region_prefix}' 开头的 processed 文件夹。")
            continue

        for j, flow_rate in enumerate(flow_rates):
            file_path = find_displacement_data_path(root_path, region_folder_name, flow_rate)

            if file_path:
                # 打印当前子图对应的数据文件夹路径
                print(f"正在处理区域: {region_prefix}, 流速: {flow_rate} m/s, 数据来源文件夹: {os.path.dirname(file_path)}")
                try:
                    df = pd.read_csv(file_path)
                    time_col, disp_col = df.columns[0], df.columns[1]
                    time = np.array(df[time_col])
                    displacement = np.array(df[disp_col])
                    fs = 1.0 / float(np.mean(np.diff(time)))
                    freqs, pxx = welch(displacement, fs, nperseg=1024)
                    pxx_db = 10 * np.log10(pxx)

                    ax.plot(freqs, pxx_db, color=colors[j % len(colors)], marker='None',
                            linestyle='-', label=f'{flow_rate}m/s')

                except Exception as e:
                    print(f"处理文件 {file_path} 时发生错误: {e}")

        # --- 对每个子图进行格式化 ---
        # 旋转Y轴标签以模仿示例样式
        ax.set_ylabel(f'截面{region_prefix.upper()}\n\n功率 (dB)', fontsize=14, rotation=0, ha='right', va='center')
        ax.yaxis.set_label_coords(-0.1, 0.5) # 调整Y轴标签位置，防止与刻度重叠

        # --- 修改点 2: 为每个子图设置X轴标签和更大的坐标刻度 ---
        ax.set_xlabel('频率 (Hz)', fontsize=14)
        ax.tick_params(axis='both', which='major', labelsize=12) # 增大坐标轴刻度字体

        ax.set_ylim(-100, -40) # 根据用户要求调整Y轴范围
        ax.set_xlim(0, 70) # 根据示例图设置X轴范围
        ax.grid(True) # 打开网格线
        # 保持上和右边框可见
        # ax.spines['top'].set_visible(False)
        # ax.spines['right'].set_visible(False)
        # 设置xy轴刻度线朝外
        ax.tick_params(axis='both', direction='out')

    # --- 对整个图表进行格式化 ---
    handles, labels = axes[0].get_legend_handles_labels()

    # --- 修改点 3: 增大图例字体大小并调整位置，使其竖向排列在图表外部右侧并添加边框 ---
    # 将图例放置在整个图表的右侧中央，并竖向排列
    fig.legend(handles, labels, loc='center right', bbox_to_anchor=(1.1, 0.5),
               fontsize=14, frameon=True) # 调整 bbox_to_anchor 的 x 值，并设置 frameon 为 True

    # --- 修改点 4: 调整子图间的垂直间距，使其更紧凑 ---
    plt.subplots_adjust(hspace=0.4) # 减小 hspace 会让子图更紧凑，增大则更松散

    # 保存图表
    output_path = os.path.join(output_folder, 'PSD_Analysis_Plot_Compact.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight') # bbox_inches='tight' 确保所有元素都被保存
    plt.close()
    print(f"PSD分析图表已成功保存至: {output_path}")


if __name__ == '__main__':
    # --- 用户配置区域 ---
    ROOT_DATA_PATH = r'J:\lqb\huanxingranliao\25CJ02'
    OUTPUT_FOLDER = 'psd_analysis_results'
    regions_to_plot = ['A', 'B', 'C']
    # 注意：根据您的示例图，流速似乎是 [0.5, 1.2, 2.0, 3.0, 3.5]
    # 我使用了您代码中定义的 [0.6, 1.2, 2.0, 3.0, 4.1, 5.0]，请根据您的实际数据调整
    flow_rates_to_plot = [0.6, 1.2, 2.0, 3.0, 4.1, 5.0]

    # --- 执行分析与绘图 ---
    plot_psd_analysis(ROOT_DATA_PATH, regions_to_plot, flow_rates_to_plot, OUTPUT_FOLDER)
